#pragma once
#include <Eigen/Core>
#include "lidar_localization/sensor_data/imu_data.hpp"
namespace lidar_localization
{
    template <typename Derived>
    static Eigen::Quaternion<typename Derived::Scalar> deltaQ(const Eigen::MatrixBase<Derived> &theta)
    {
        typedef typename Derived::Scalar Scalar_t;

        Eigen::Quaternion<Scalar_t> dq;
        Eigen::Matrix<Scalar_t, 3, 1> half_theta = theta;
        half_theta /= static_cast<Scalar_t>(2.0);
        dq.w() = static_cast<Scalar_t>(1.0);
        dq.x() = half_theta.x();
        dq.y() = half_theta.y();
        dq.z() = half_theta.z();
        // dq.normalize();
        return dq;
    }
    template <typename Derived>
    static Eigen::Matrix<typename Derived::Scalar, 3, 3> skewSymmetric(const Eigen::MatrixBase<Derived> &q)
    {
        Eigen::Matrix<typename Derived::Scalar, 3, 3> ans;
        ans << typename Derived::Scalar(0), -q(2), q(1),
            q(2), typename Derived::Scalar(0), -q(0),
            -q(1), q(0), typename Derived::Scalar(0);
        return ans;
    }

    template <typename Derived>
    static Eigen::Quaternion<typename Derived::Scalar> positify(const Eigen::QuaternionBase<Derived> &q)
    {
        return q;
    }
    template <typename Derived>
    static Eigen::Matrix<typename Derived::Scalar, 4, 4> Qleft(const Eigen::QuaternionBase<Derived> &q)
    {
        Eigen::Quaternion<typename Derived::Scalar> qq = positify(q);
        Eigen::Matrix<typename Derived::Scalar, 4, 4> ans;
        ans(0, 0) = qq.w(), ans.template block<1, 3>(0, 1) = -qq.vec().transpose();
        ans.template block<3, 1>(1, 0) = qq.vec(), ans.template block<3, 3>(1, 1) = qq.w() * Eigen::Matrix<typename Derived::Scalar, 3, 3>::Identity() + skewSymmetric(qq.vec());
        return ans;
    }

    template <typename Derived>
    static Eigen::Matrix<typename Derived::Scalar, 4, 4> Qright(const Eigen::QuaternionBase<Derived> &p)
    {
        Eigen::Quaternion<typename Derived::Scalar> pp = positify(p);
        Eigen::Matrix<typename Derived::Scalar, 4, 4> ans;
        ans(0, 0) = pp.w(), ans.template block<1, 3>(0, 1) = -pp.vec().transpose();
        ans.template block<3, 1>(1, 0) = pp.vec(), ans.template block<3, 3>(1, 1) = pp.w() * Eigen::Matrix<typename Derived::Scalar, 3, 3>::Identity() - skewSymmetric(pp.vec());
        return ans;
    }

    class PreIntegration
    {
    public:
        PreIntegration() = delete;
        PreIntegration(const std::deque<IMUData> &imu_data, const Eigen::Vector3d &linearized_ba, const Eigen::Vector3d &linearized_bg, const Eigen::VectorXd &noise, const Eigen::Vector3d &gravity)
            : imu_data_(imu_data),
              linearized_ba_(linearized_ba),
              linearized_bg_(linearized_bg),
              jacobian_(Eigen::Matrix<double, 15, 15>::Identity()),
              covariance_(Eigen::Matrix<double, 15, 15>::Zero()),
              gravity_(gravity)
        {
            sumdt_ = imu_data_.back().time - imu_data_.front().time;
            noise_ = Eigen::Matrix<double, 18, 18>::Zero();
            for (int i = 0; i < 3; ++i)
            {
                noise_(i, i) = noise(i) * noise(i);
                noise_(i + 3, i + 3) = noise(i + 3) * noise(i + 3);
                noise_(i + 6, i + 6) = noise(i) * noise(i);
                noise_(i + 9, i + 9) = noise(i + 3) * noise(i + 3);
                noise_(i + 12, i + 12) = noise(i + 6) * noise(i + 6);
                noise_(i + 15, i + 15) = noise(i + 9) * noise(i + 9);
            }
            alpha_.setZero();
            dq_.setIdentity();
            beta_.setZero();
            propagate();
        }

        void repropagate(const Eigen::Vector3d &linearized_ba, const Eigen::Vector3d &linearized_bg)
        {
            alpha_.setZero();
            dq_.setIdentity();
            beta_.setZero();
            linearized_ba_ = linearized_ba;
            linearized_bg_ = linearized_bg;
            jacobian_.setIdentity();
            covariance_.setZero();
            propagate();
        }

        void midPointIntegration(double dt,
                                 const Eigen::Vector3d &acc0,
                                 const Eigen::Vector3d &gyr0,
                                 const Eigen::Vector3d &acc1,
                                 const Eigen::Vector3d &gyr1,
                                 const Eigen::Vector3d &alpha,
                                 const Eigen::Quaterniond &dq,
                                 const Eigen::Vector3d &beta,
                                 const Eigen::Vector3d &linearized_ba,
                                 const Eigen::Vector3d &linearized_bg,
                                 Eigen::Vector3d &result_alpha,
                                 Eigen::Quaterniond &result_dq,
                                 Eigen::Vector3d &result_beta,
                                 Eigen::Vector3d &result_linearized_ba,
                                 Eigen::Vector3d &result_linearized_bg,
                                 bool update_jacobian)
        {
            Eigen::Vector3d gyro_av = 0.5 * (gyr0 + gyr1) - linearized_bg;
            Eigen::Quaterniond temp = Eigen::Quaterniond(1, gyro_av(0) * dt / 2, gyro_av(1) * dt / 2, gyro_av(2) * dt / 2);
            temp.normalize();
            result_dq = dq * temp;
            Eigen::Vector3d acc_av = 0.5 * ((dq * (acc0 - linearized_ba) + result_dq * (acc1 - linearized_ba)));
            result_alpha = alpha + beta * dt + 0.5 * acc_av * dt * dt;
            result_beta = beta + acc_av * dt;
            result_linearized_ba = linearized_ba;
            result_linearized_bg = linearized_bg;

            if (update_jacobian)
            {
                Eigen::Vector3d a0_x = acc0 - linearized_ba;
                Eigen::Vector3d a1_x = acc1 - linearized_ba;
                Eigen::Matrix3d R_w_x, R_a0_x, R_a1_x;

                R_w_x << 0, -gyro_av(2), gyro_av(1),
                    gyro_av(2), 0, -gyro_av(0),
                    -gyro_av(1), gyro_av(0), 0;
                R_a0_x << 0, -a0_x(2), a0_x(1),
                    a0_x(2), 0, -a0_x(0),
                    -a0_x(1), a0_x(0), 0;
                R_a1_x << 0, -a1_x(2), a1_x(1),
                    a1_x(2), 0, -a1_x(0),
                    -a1_x(1), a1_x(0), 0;

                Eigen::Matrix<double, 15, 15> F = Eigen::Matrix<double, 15, 15>::Zero();
                F.block<3, 3>(0, 0) = Eigen::Matrix3d::Identity();
                F.block<3, 3>(0, 3) = -0.25 * dt * dt * (dq.toRotationMatrix() * R_a0_x + result_dq.toRotationMatrix() * R_a1_x * (Eigen::Matrix3d::Identity() - R_w_x * dt));
                F.block<3, 3>(0, 6) = Eigen::Matrix3d::Identity() * dt;
                F.block<3, 3>(0, 9) = -0.25 * dt * dt * (dq.toRotationMatrix() + result_dq.toRotationMatrix());
                F.block<3, 3>(0, 12) = 0.25 * dt * dt * dt * result_dq.toRotationMatrix() * R_a1_x;

                F.block<3, 3>(3, 3) = Eigen::Matrix3d::Identity() - R_w_x * dt;
                F.block<3, 3>(3, 12) = -1.0 * dt * Eigen::Matrix3d::Identity(3, 3);

                F.block<3, 3>(6, 3) = -0.5 * dt * (dq.toRotationMatrix() * R_a0_x + result_dq.toRotationMatrix() * R_a1_x * (Eigen::Matrix3d::Identity() - R_w_x * dt));
                F.block<3, 3>(6, 6) = Eigen::Matrix3d::Identity();
                F.block<3, 3>(6, 9) = -0.5 * dt * (dq.toRotationMatrix() + result_dq.toRotationMatrix());
                F.block<3, 3>(6, 12) = 0.5 * dt * dt * result_dq.toRotationMatrix() * R_a1_x;

                F.block<3, 3>(9, 9) = Eigen::Matrix3d::Identity();

                F.block<3, 3>(12, 12) = Eigen::Matrix3d::Identity();

                Eigen::Matrix<double, 15, 18> G = Eigen::Matrix<double, 15, 18>::Zero();
                G.block<3, 3>(0, 0) = 0.25 * dt * dt * dq.toRotationMatrix();
                G.block<3, 3>(0, 3) = -0.125 * dt * dt * dt * result_dq.toRotationMatrix() * R_a1_x;
                G.block<3, 3>(0, 6) = 0.25 * dt * dt * result_dq.toRotationMatrix();
                G.block<3, 3>(0, 9) = G.block<3, 3>(0, 3);

                G.block<3, 3>(3, 3) = 0.5 * dt * Eigen::Matrix3d::Identity(3, 3);
                G.block<3, 3>(3, 9) = 0.5 * dt * Eigen::Matrix3d::Identity(3, 3);

                G.block<3, 3>(6, 0) = 0.5 * dt * dq.toRotationMatrix();
                G.block<3, 3>(6, 3) = -0.25 * dt * dt * result_dq.toRotationMatrix() * R_a1_x;
                G.block<3, 3>(6, 6) = 0.5 * dt * result_dq.toRotationMatrix();
                G.block<3, 3>(6, 9) = G.block<3, 3>(6, 3);

                G.block<3, 3>(9, 12) = Eigen::Matrix3d::Identity() * dt;

                G.block<3, 3>(12, 15) = Eigen::Matrix3d::Identity() * dt;

                jacobian_ = F * jacobian_;
                covariance_ = F * covariance_ * F.transpose() + G * noise_ * G.transpose();
            }
        }

        void propagate()
        {
            for (int i = 1; i < static_cast<int>(imu_data_.size()); ++i)
            {
                double dt = imu_data_[i].time - imu_data_[i - 1].time;

                Eigen::Vector3d acc0;
                acc0[0] = imu_data_[i - 1].linear_acceleration.x;
                acc0[1] = imu_data_[i - 1].linear_acceleration.y;
                acc0[2] = imu_data_[i - 1].linear_acceleration.z;

                Eigen::Vector3d gry0;
                gry0[0] = imu_data_[i - 1].angular_velocity.x;
                gry0[1] = imu_data_[i - 1].angular_velocity.y;
                gry0[2] = imu_data_[i - 1].angular_velocity.z;

                Eigen::Vector3d acc1;
                acc1[0] = imu_data_[i].linear_acceleration.x;
                acc1[1] = imu_data_[i].linear_acceleration.y;
                acc1[2] = imu_data_[i].linear_acceleration.z;

                Eigen::Vector3d gry1;
                gry1[0] = imu_data_[i].angular_velocity.x;
                gry1[1] = imu_data_[i].angular_velocity.y;
                gry1[2] = imu_data_[i].angular_velocity.z;

                Eigen::Vector3d result_alpha;
                Eigen::Quaterniond result_dq;
                Eigen::Vector3d result_beta;
                Eigen::Vector3d result_linearized_ba;
                Eigen::Vector3d result_linearized_bg;
                midPointIntegration(dt, acc0, gry0, acc1, gry1,
                                    alpha_, dq_, beta_, linearized_ba_,
                                    linearized_bg_, result_alpha, result_dq, result_beta, result_linearized_ba, result_linearized_bg, true);
                alpha_ = result_alpha;
                dq_ = result_dq;
                beta_ = result_beta;
                linearized_ba_ = result_linearized_ba;
                linearized_bg_ = result_linearized_bg;
                dq_.normalize();
            }
        }

        Eigen::Matrix<double, 15, 1> evaluate(const Eigen::Vector3d &Pi,
                                              const Eigen::Quaterniond &Qi,
                                              const Eigen::Vector3d &Vi,
                                              const Eigen::Vector3d &Bai,
                                              const Eigen::Vector3d &Bgi,
                                              const Eigen::Vector3d &Pj,
                                              const Eigen::Quaterniond &Qj,
                                              const Eigen::Vector3d &Vj,
                                              const Eigen::Vector3d &Baj,
                                              const Eigen::Vector3d &Bgj)
        {
            Eigen::Matrix<double, 15, 1> residuals;

            Eigen::Vector3d dba = Bai - linearized_ba_;
            Eigen::Vector3d dbg = Bgi - linearized_bg_;
            Eigen::Matrix3d dalpha_dba = jacobian_.block<3, 3>(0, 9);
            // Eigen::Matrix3d dq_dba = jacobian_.block<3, 3>(3, 9);
            Eigen::Matrix3d dbeta_dba = jacobian_.block<3, 3>(6, 9);
            Eigen::Matrix3d dalpha_dbg = jacobian_.block<3, 3>(0, 12);
            Eigen::Matrix3d dq_dbg = jacobian_.block<3, 3>(3, 12);
            Eigen::Matrix3d dbeta_dbg = jacobian_.block<3, 3>(6, 12);

            Eigen::Quaterniond corrected_dq = dq_ * deltaQ(dq_dbg * dbg);
            Eigen::Vector3d corrected_alpha = alpha_ + dalpha_dba * dba + dalpha_dbg * dbg;
            Eigen::Vector3d corrected_beta = beta_ + dbeta_dba * dba + dbeta_dbg * dbg;

            residuals.block<3, 1>(0, 0) = Qi.inverse() * (Pj - Pi - Vi * sumdt_ - 0.5 * gravity_ * sumdt_ * sumdt_) - corrected_alpha;
            residuals.block<3, 1>(3, 0) = 2 * (corrected_dq.inverse() * (Qi.inverse() * Qj)).vec();
            residuals.block<3, 1>(6, 0) = Qi.inverse() * (Vj - Vi - gravity_ * sumdt_) - corrected_beta;
            residuals.block<3, 1>(9, 0) = Baj - Bai;
            residuals.block<3, 1>(12, 0) = Bgj - Bgi;
            return residuals;
        }

        Eigen::Matrix<double, 15, 30> jacobian(const Eigen::Vector3d &Pi,
                                               const Eigen::Quaterniond &Qi,
                                               const Eigen::Vector3d &Vi,
                                               const Eigen::Vector3d &Bai,
                                               const Eigen::Vector3d &Bgi,
                                               const Eigen::Vector3d &Pj,
                                               const Eigen::Quaterniond &Qj,
                                               const Eigen::Vector3d &Vj,
                                               const Eigen::Vector3d &Baj,
                                               const Eigen::Vector3d &Bgj)
        {
            Eigen::Matrix<double, 15, 30> jacobians;
            jacobians.setZero();

            // Eigen::Vector3d dba = Bai - linearized_ba_;
            Eigen::Vector3d dbg = Bgi - linearized_bg_;
            Eigen::Matrix3d dalpha_dba = jacobian_.block<3, 3>(0, 9);
            // Eigen::Matrix3d dq_dba = jacobian_.block<3, 3>(3, 9);
            Eigen::Matrix3d dbeta_dba = jacobian_.block<3, 3>(6, 9);
            Eigen::Matrix3d dalpha_dbg = jacobian_.block<3, 3>(0, 12);
            Eigen::Matrix3d dq_dbg = jacobian_.block<3, 3>(3, 12);
            Eigen::Matrix3d dbeta_dbg = jacobian_.block<3, 3>(6, 12);
            Eigen::Quaterniond corrected_dq = dq_ * deltaQ(dq_dbg * dbg);
            // Eigen::Vector3d corrected_alpha = alpha_ + dalpha_dba * dba + dalpha_dbg * dbg;
            // Eigen::Vector3d corrected_beta = beta_ + dbeta_dba * dba + dbeta_dbg * dbg;

            jacobians.block<3, 3>(0, 0) = -Qi.inverse().toRotationMatrix();
            jacobians.block<3, 3>(0, 3) = skewSymmetric(Qi.inverse() * (Pj - Pi - Vi * sumdt_ - 0.5 * gravity_ * sumdt_ * sumdt_));
            jacobians.block<3, 3>(0, 6) = -Qi.inverse().toRotationMatrix() * sumdt_;
            jacobians.block<3, 3>(0, 9) = -dalpha_dba;
            jacobians.block<3, 3>(0, 12) = -dalpha_dbg;

            jacobians.block<3, 3>(0, 15) = Qi.inverse().toRotationMatrix();

            jacobians.block<3, 3>(3, 3) = -(Qleft(Qj.inverse() * Qi) * Qright(corrected_dq)).bottomRightCorner<3, 3>();
            jacobians.block<3, 3>(3, 12) = -(Qleft(Qj.inverse() * Qi * corrected_dq)).bottomRightCorner<3, 3>() * dq_dbg;

            jacobians.block<3, 3>(3, 18) = (Qleft(corrected_dq.inverse() * Qi.inverse() * Qj)).bottomRightCorner<3, 3>();

            jacobians.block<3, 3>(6, 3) = skewSymmetric(Qi.inverse() * (Vj - Vi - gravity_ * sumdt_));
            jacobians.block<3, 3>(6, 6) = -Qi.inverse().toRotationMatrix();
            jacobians.block<3, 3>(6, 9) = -dbeta_dba;
            jacobians.block<3, 3>(6, 12) = -dbeta_dbg;

            jacobians.block<3, 3>(6, 21) = Qi.inverse().toRotationMatrix();

            jacobians.block<3, 3>(9, 9) = -Eigen::Matrix3d::Identity();

            jacobians.block<3, 3>(9, 24) = Eigen::Matrix3d::Identity();

            jacobians.block<3, 3>(12, 12) = -Eigen::Matrix3d::Identity();

            jacobians.block<3, 3>(12, 27) = Eigen::Matrix3d::Identity();

            return jacobians;
        }

        std::deque<IMUData> imu_data_;
        double sumdt_;
        Eigen::Vector3d linearized_ba_;
        Eigen::Vector3d linearized_bg_;

        Eigen::Matrix<double, 15, 15> jacobian_, covariance_;
        Eigen::Matrix<double, 18, 18> noise_;
        Eigen::Vector3d gravity_;

        Eigen::Vector3d alpha_;
        Eigen::Vector3d beta_;
        Eigen::Quaterniond dq_;
    };
} // namespace lidar_localization